#!/usr/bin/env python3
"""
简单UI测试节点 - 验证正确的UI返回格式
"""

from .lora_trainer_utils.constants import BASIC_LORA_TRAINING_CATEGORY

class SimpleUITestNode:
    """简单UI测试节点"""
    
    aux_id = "comfyui_lora_train"
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "test_type": (["成功", "错误", "进度"], {
                    "default": "成功",
                    "tooltip": "选择测试类型"
                }),
                "message": ("STRING", {
                    "default": "测试消息",
                    "tooltip": "测试消息"
                }),
            }
        }
    
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("结果",)
    FUNCTION = "test_ui"
    CATEGORY = BASIC_LORA_TRAINING_CATEGORY
    
    def test_ui(self, test_type, message):
        """测试UI返回格式"""
        
        if test_type == "成功":
            return {
                "ui": {
                    "status": ["✅ 成功"],
                    "message": [message]
                },
                "result": (f"成功: {message}",)
            }
            
        elif test_type == "错误":
            return {
                "ui": {
                    "error": [f"❌ 错误: {message}"]
                },
                "result": (f"错误: {message}",)
            }
            
        elif test_type == "进度":
            return {
                "ui": {
                    "status": ["🔄 处理中"],
                    "progress": [50]
                },
                "result": (f"进度: {message}",)
            }
        
        return {
            "ui": {},
            "result": (f"未知类型: {message}",)
        }


# 注册节点
NODE_CLASS_MAPPINGS = {
    "SimpleUITestNode": SimpleUITestNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SimpleUITestNode": "🧪 简单UI测试",
} 